import streamlit as st
import requests
import time

# Set the app to wide mode
st.set_page_config(layout="wide")

# Function to interact with the LLM API and measure response time
def ask_question(api_url, model, question):
    headers = {
        "Content-Type": "application/json",
    }

    payload = {
        "model": model,  # Selected LLM model
        "prompt": question,
        "max_tokens": 1000  # Adjust as needed
    }

    start_time = time.time()  # Record start time
    response = requests.post(api_url, headers=headers, json=payload)
    end_time = time.time()  # Record end time

    response_time = end_time - start_time  # Calculate response time

    if response.status_code == 200:
        answer = response.json()["choices"][0]["text"]
        return answer, response_time
    else:
        return "错误: 无法从LLM获取响应。", response_time  # Error message in Chinese

# Function to display text, response time, and copy button in separate sections
def display_text_with_copy(text, label, height, response_time=None):
    st.text_area(label, value=text, height=height, max_chars=None, key=f"text_{label}", disabled=True)

    if response_time is not None:
        st.markdown(f"**响应时间:** {response_time:.2f} 秒")  # Display response time below the text area

    st.button(f"复制 {label}", key=f"copy_{label}")  # Display copy button separately

# Streamlit UI
st.title("Falcomm 本地 LLM 问答测试系统")  # Local LLM Q&A Demo

# Input for LLM API URL
api_url = st.text_input("请输入 LLM API URL:", value="http://localhost:8000/v1/completions")

# Model selection
model = st.selectbox(
    "选择 LLM 模型:",  # Select LLM Model
    options=["qwen2-0.5b", "qwen2-1.5b", "other-model-1", "other-model-2"]
)

# Set the column widths to a 3:2 ratio
col1, col2 = st.columns([3, 2])  # 3:2 ratio for the columns

with col1:
    # Input question and display the answer in the left column
    with st.form(key="question_form"):
        question = st.text_input("请输入你的问题:")  # Please enter your question
        submit_button = st.form_submit_button("提交")  # Submit

    if submit_button and question:
        # Get the answer from the LLM and measure response time
        answer, response_time = ask_question(api_url, model, question)

        # Append the question, answer, and response time to the history
        st.session_state.history.append((question, answer, response_time))

        # Display the answer with auto-wrap, response time, and copy button
        display_text_with_copy(answer, "回答", height=600, response_time=response_time)  # Answer with larger height
    elif not question and submit_button:
        st.write("请输入一个问题。")  # Please enter a question

with col2:
    # Display the conversation history in the right column with each Q&A pair collapsible
    if "history" not in st.session_state:
        st.session_state.history = []

    st.subheader("对话历史")  # Conversation History
    for idx, (q, a, rt) in enumerate(reversed(st.session_state.history), 1):
        with st.expander(f"问题 {idx}: {q}"):  # Question {idx}: {q}
            display_text_with_copy(q, f"问题 {idx}", height=100)  # Copy Question with smaller height
            display_text_with_copy(a, f"回答 {idx}", height=600, response_time=rt)  # Copy Answer with larger height and response time
